﻿//    TestGenerator - experiment with piping data into ffmpeg in "pipefmt"
//    Copyright (C) 2012  Juergen Bohl

//    This program is free software: you can redistribute it and/or modify
//    it under the terms of the GNU General Public License as published by
//    the Free Software Foundation, either version 3 of the License, or
//    (at your option) any later version.

//    This program is distributed in the hope that it will be useful,
//    but WITHOUT ANY WARRANTY; without even the implied warranty of
//    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//    GNU General Public License for more details.

//    You should have received a copy of the GNU General Public License
//    along with this program.  If not, see <http://www.gnu.org/licenses/>.

namespace TestGenerator
{
    #region using directives
    using System;
    using System.Collections.Generic;
    using System.IO;
    using System.Linq;
    using System.Runtime.InteropServices;
    using System.Text;
    #endregion

    partial class PipeFmtFileOutput : IAVSink, IDisposable
    {
        private List<VideoFormat> videoStreams;

        private List<AudioFormat> audioStreams;

        private FileStream fileStream;

        private long imageCounter;

        private int streamCounter;

        private int VideoStreamsCount
        {
            get { return (this.videoStreams != null) ? this.videoStreams.Count : 0; }
        }

        private int AudioStreamsCount
        {
            get { return (this.audioStreams != null) ? this.audioStreams.Count : 0; }
        }

        /// <summary>Gets a value indicating whether the next packet we expect is a video stream packet. </summary>
        ///
        /// <value>True if next expected packet is video stream packet, false if not. </value>
        private bool NextIsVideoStream
        {
            get
            {
                if (this.streamCounter >= this.VideoStreamsCount)
                    return false;

                return true;
            }
        }

        private bool NextIsAudioStream
        {
            get { return !this.NextIsVideoStream; }
        }

        private int NextStreamNumber
        {
            get
            {
                if (this.NextIsVideoStream)
                {
                    return this.streamCounter;
                }
                else
                {
                    return this.streamCounter - this.VideoStreamsCount;
                }
            }
        }

        public void Initialize(InitOptions initOptions)
        {
            this.videoStreams = initOptions.videoStreams != null ? initOptions.videoStreams.ToList() : null;
            this.audioStreams = initOptions.audioStreams != null ? initOptions.audioStreams.ToList() : null;

            this.imageCounter = 0;
            this.streamCounter = 0;

            this.fileStream = new FileStream(initOptions.fileNameOutput, FileMode.CreateNew);
            this.WriteHeader();
        }

        public void Dispose()
        {
            this.Dispose(true);
            GC.SuppressFinalize(this);
        }

        /// <summary>
        /// Releases unmanaged and - optionally - managed resources.
        /// </summary>
        /// <param name="disposing"><c>true</c> to release both managed and unmanaged resources; <c>false</c> to release only unmanaged resources.</param>
        protected virtual void Dispose(bool disposing)
        {
            if (disposing)
            {
                // Release managed resources:
                // Call Dispose() on other objects owned by this instance.
                if (this.fileStream != null)
                {
                    this.fileStream.Dispose();
                }
            }

            // Release unmanaged resources owned by (just) this instance.
        }

        public struct InitOptions
        {
            public InitOptions(string filename, IEnumerable<VideoFormat> videoStreams, IEnumerable<AudioFormat> audioStreams)
            {
                this.fileNameOutput = filename;
                this.videoStreams = videoStreams;
                this.audioStreams = audioStreams;
            }

            public InitOptions(string filename, VideoFormat videoFormat, AudioFormat audioFormat)
                : this(filename, videoFormat != null ? Enumerable.Repeat(videoFormat, 1) : null, audioFormat != null ? Enumerable.Repeat(audioFormat, 1) : null)
            {
            }

            public string fileNameOutput;

            public IEnumerable<VideoFormat> videoStreams;

            public IEnumerable<AudioFormat> audioStreams;
        }

        private void WriteHeader()
        {
            PipeFormatHelper.WriteHeader(
                                    buffer =>
                                    {
                                        this.fileStream.Write(buffer, 0, buffer.Length);
                                        return true;
                                    },
                                        videoStreams,
                                        audioStreams);
        }
    }

    partial class PipeFmtFileOutput
    {
        public void AddFrameRgb24(IntPtr bitmap, int size, int stride)
        {
            if (!this.NextIsVideoStream)
            {
                throw new InvalidOperationException("Expected an audio packet at this time.");
            }

            int streamNo = this.NextStreamNumber;
            VideoFormat vf = this.videoStreams[streamNo];
            this.WriteFrameHeaderVideo(streamNo);
            byte[] buffer = this.ConvertToYv12Bitmap(bitmap, vf.Width, vf.Height, stride);
            this.fileStream.Write(buffer, 0, buffer.Count());
            this.Increment();
        }

        public void AddAudioSamplesInterleaved(short[] samples)
        {
            if (!this.NextIsAudioStream)
            {
                throw new InvalidOperationException("Expected an audio packet at this time.");
            }

            int streamNo = this.NextStreamNumber;
            AudioFormat af = this.audioStreams[streamNo];
            this.WriteFrameHeaderAudio(0, samples.Length / 2);
            byte[] byteBuffer = new byte[samples.Length * sizeof(short)];
            Buffer.BlockCopy(
                            samples,
                            0,
                            byteBuffer,
                            0,
                            samples.Length * sizeof(short));
            this.fileStream.Write(byteBuffer, 0, byteBuffer.Length);
            this.Increment();
        }

        private void Increment()
        {
            this.streamCounter++;
            if (this.streamCounter >= (this.VideoStreamsCount + this.AudioStreamsCount))
            {
                this.streamCounter = 0;
                this.imageCounter++;
            }
        }

        private void WriteFrameHeaderAudio(int streamNo, int samplesCount)
        {
            string s = string.Format("ASTREAM#{0}: C={1}\n", streamNo, samplesCount);
            byte[] strAscii = Encoding.ASCII.GetBytes(s);
            this.fileStream.Write(strAscii, 0, strAscii.Count());
        }

        private byte[] ConvertToYv12Bitmap(IntPtr bitmapRgb24, int width, int height, int stride)
        {
            byte[] yv12Buffer = new byte[3 * width * height / 2];
            GCHandle pinnedDest = GCHandle.Alloc(yv12Buffer, GCHandleType.Pinned);
            BGRtoYV12.ConvertRgb24ToYv12(
                                    bitmapRgb24,
                                    stride,
                                    width,
                                    height,
                                    pinnedDest.AddrOfPinnedObject(),
                                    width,
                                    pinnedDest.AddrOfPinnedObject() + width * height,
                                    pinnedDest.AddrOfPinnedObject() + 5 * width * height / 4,
                                    width / 2);
            pinnedDest.Free();
            return yv12Buffer;
        }

        private void WriteFrameHeaderVideo(int streamNo)
        {
            string s = string.Format("VSTREAM#{0}:\n", streamNo);
            byte[] strAscii = Encoding.ASCII.GetBytes(s);
            this.fileStream.Write(strAscii, 0, strAscii.Count());
        }
    }
}
